package pd_device

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"log"
	"math/big"
	"time"

	"gitee.com/general252/proxy_device/pd_common"
	"github.com/armon/go-socks5"
	"github.com/gorilla/websocket"
	"github.com/quic-go/quic-go"
)

type Device struct {
	wsURL string
}

func NewDevice(wsURL string) *Device {
	return &Device{
		wsURL: wsURL,
	}
}

func (tis *Device) Connect() error {
	dialer := &websocket.Dialer{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}

	// 连接websocket
	url := tis.wsURL
	conn, _, err := dialer.Dial(url, nil)
	if err != nil {
		log.Println(err, url)
		return err
	}
	log.Printf("dial %v success. 连接websocket服务器成功", url)

	// quic server
	listener, err := quic.Listen(pd_common.NewWebSocketPacketConn(conn), generateTLSConfig(), &quic.Config{
		KeepAlivePeriod: time.Second * 10,
	})
	if err != nil {
		log.Println(err)
		return err
	}
	defer listener.Close()
	log.Println("quic listen success. 启动quic服务器成功")

	// socks5 server
	socks5Proxy, err := socks5.New(&socks5.Config{})
	if err != nil {
		log.Println()
		return err
	}
	log.Println("socks5 server success. 启动socks5服务器成功")

	for {
		quicConn, err := listener.Accept(context.TODO())
		if err != nil {
			log.Println(err)
			break
		}
		log.Println("new quic conn. 收到新的quic连接")

		go func(quicConn quic.Connection) {
			defer func() {
				_ = quicConn.CloseWithError(quic.ApplicationErrorCode(256), "no error")
				log.Println("quic conn closed. 关闭quic连接")
			}()

			for {
				// quic stream
				stream, err := quicConn.AcceptStream(context.TODO())
				if err != nil {
					log.Println(err)
					return
				}
				log.Printf("new quic stream: %v. 收到quic流", stream.StreamID())

				go func(stream quic.Stream) {
					defer func() {
						_ = stream.Close()
						log.Printf("quic stream %v closed. 关闭quic流", stream.StreamID())
					}()

					// 处理socks5数据
					_ = socks5Proxy.ServeConn(NewQuicStream(stream))
				}(stream)
			}
		}(quicConn)
	}

	return nil
}

// a bare-bones TLS config for the server
func generateTLSConfig() *tls.Config {
	// _, priv, err := ed25519.GenerateKey(rand.Reader)
	priv, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		panic(err)
	}
	template := x509.Certificate{SerialNumber: big.NewInt(1)}
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, priv.Public(), priv)
	if err != nil {
		panic(err)
	}

	return &tls.Config{
		Certificates: []tls.Certificate{{
			Certificate: [][]byte{certDER},
			PrivateKey:  priv,
		}},
		NextProtos: []string{"quic-echo-example"},
	}
}
